import pickle

import torch

import chatbot.lib as lib
from chatbot.classify.classify import Classify
from chatbot.dnn.sort import word_sequence
from chatbot.prepar_corpus.prepar_dnn_corpus import recall_corpus
from classify.build_model import get_classify_model
from chatbot_core.word_sequence import WordSequence
import chatbot.config as config
import chatbot_core.dataset as dataset
import chatbot.chatbot_core.seq2seq as seq2seq
from chatbot.dnn.recall.recall import Recall
from chatbot.dnn.sort.dnnsort import DnnSort

import fasttext

fasttext.FastText.eprint = lambda x: None


def test_classify_model(by_word):
    model = get_classify_model(by_word)
    text = ["吃 了 吗？", "java", "devops 使用 argocd 合适 吗"]
    ret = model.predict(text)
    print(ret)


def create_word_sequence_model():
    ws = WordSequence()
    for line in open(config.chatbot_input_path, encoding="utf-8").readlines():
        # 输入
        ws.fit(line.strip().split())
    ws.build_vocab()
    print(len(ws))
    pickle.dump(ws, open(config.chatbot_ws_input_path, "wb"))

    ws = WordSequence()
    for line in open(config.chatbot_target_path, encoding="utf-8").readlines():
        # 输入
        ws.fit(line.strip().split())
    ws.build_vocab()
    print(len(ws))
    pickle.dump(ws, open(config.chatbot_ws_target_path, "wb"))


def test_dataset():
    for input, label, input_length, label_length in dataset.train_data_loader:
        print(input)
        print(label)
        print(input_length)
        print(label_length)
        break


def test_dnn(by_word, method):
    recall = Recall(by_word=by_word, method=method)
    sentence = {"cut_by_word": "python 是 什 么", "cut": "python 是 什么", "entity": ["python"]}
    results = recall.predict(sentence)
    print(results)


def test_dnn_dataset(by_word=True):
    import chatbot.dnn.sort.dataset as ds
    for input, label, input_length, label_length in ds.dnnsort_data_loader:
        print(input)
        print(label)
        print(input_length)
        print(label_length)
        break


def dnn_sort_train(epoch):
    import chatbot.dnn.sort.train as train
    for i in range(epoch):
        train.train(i)


def dnn_sort_eval():
    import chatbot.dnn.sort.eval as eval
    eval.eval(True)


def dnn_sort_class_test():
    ds = DnnSort()
    recall = Recall(by_word=True, method="fasttext")
    sentence = {"cut_by_word": "python 是 什 么", "cut": "python 是 什么", "entity": ["python"]}
    recall_list = recall.predict(sentence)
    answer = ds.predict(sentence, recall_list)
    print(answer)


class ChatBot:
    def __init__(self):
        self.classify = Classify()
        self.seq2seq = seq2seq.Seq2seq().to(config.device)
        self.seq2seq.load_state_dict(torch.load(config.chatbot_model_save_path))
        self.dnnSort = DnnSort()
        self.recall = Recall(by_word=True, method="fasttext")

    def predict(self, sentence):
        sentence = {"cut": " ".join(lib.cut(sentence, by_word=False)),
                    "cut_by_word": " ".join(lib.cut(sentence, by_word=True)), "entity": lib.get_entity(sentence)}
        if self.isQA(sentence):
            recall_list = self.recall.predict(sentence)
            answer = self.dnnSort.predict(sentence, recall_list)
            return "问答:" + answer
        else:
            sentence = sentence["cut_by_word"] if by_word else sentence["cut"]
            sentence = sentence.split(" ")
            input_length = torch.LongTensor(
                [len(sentence) if len(sentence) > config.chatbot_input_max_len else config.chatbot_input_max_len]).to(
                config.device)
            question = torch.LongTensor(
                [config.chatbot_ws_input.transform(sentence, max_len=config.chatbot_input_max_len)]).to(config.device)
            indices = self.seq2seq.evaluation(question, input_length)
            return "闲聊：" + "".join(config.chatbot_ws_target.inverse_transform(indices))

    def isQA(self, sentence):
        ret, _ = self.classify.predict(sentence)
        if ret == "QA":
            return True
        return False


def run():
    cb = ChatBot()
    while True:
        q = input("请说：")
        if q == "exit":
            exit()
        a = cb.predict(q)
        print(a)

if __name__ == '__main__':
    # test_user_dict.test4()
    # test_user_dict.test4(True)

    by_word = True
    # 先处理语料，生成文件
    # process(by_word)

    # 创建fasttext模型
    # build_classify_model(by_word)

    # acc=eval(by_word)
    # print(acc)

    # test_classify_model(by_word)

    # classify = Classify()
    # dict = {"cut": "java", "cut_by_word": "java"}
    # ret = classify.predict(dict)
    # print(ret)

    # 生成问答语料
    # input_path = "corpus/chatbot/input.txt"
    # input_by_word_path = "corpus/chatbot/input_by_word.txt"
    # target_path = "corpus/chatbot/target.txt"
    # target_by_word_path = "corpus/chatbot/target_by_word.txt"
    # chatbot_corpus.prepar_xiaohuangji(by_word=False)
    # os.remove(input_by_word_path)
    # os.remove(target_by_word_path)
    # chatbot_corpus.prepar_xiaohuangji(by_word=True)

    # create_word_sequence_model()

    # test_dataset()

    # for i in range(1):
    #     train.train(i)

    # eval.eval()

    # recall_corpus.prepar_recall_corpus()

    # test_dnn(by_word,"bm25")
    # test_dnn(by_word,"fasttext")

    # fasttext_vectors.build_model()

    # word_sequence.save_dnnsort_ws(True)
    # word_sequence.save_dnnsort_ws(False)

    # test_dnn_dataset()

    # dnn_sort_train(1)
    # dnn_sort_eval()

    # dnn_sort_class_test()

    run()